# Make graphs for figure 1

# Load Models and Store AUCs
countries <- c('indo', 'colo')

symSignRescale <- function(x) {
  C = .5
  return(sign(x)*log(1+abs(x)/10^C))
}
ticks <- c(-200, -50,-15,-5, -2, 0, 2, 5, 15)
xtics <- list(indo = list(ticks = symSignRescale(ticks),
                          labels = ticks),
              colo = list(ticks = -5:5,
                          labels = -5:5))
for (country in countries) {
    source(paste("config_code/",country,"_data_setup.R",sep = ""))
    source(paste("config_code/",country,"_predictor_vars.R",sep = ""))
    models <- names(figures_1_2)[-1]
    T <- end.year - start.year + 1
    
    # Get population benchmark dataset
    pop_auc_array <- array(NA,
                           dim = T)
    filename <- paste(modeldir,"/",
                      country,
                      "_ebma_",
                      "count",
                      "_population.RData",
                      sep = "")
    load(filename)
    for (i in 1:T) {
      predictions <- as.vector(ebma.results[[i]]$fit.oos)
      realizations <- as.vector(ebma.results[[i]]$actual.oos)
      predictions[predictions > max(realizations)] <- max(realizations)
      pop_auc_array[i] <- mean((realizations - predictions)^2)
    }
    
    auc_array <- array(NA,
                       dim = c(length(models),
                               T),
                       dimnames = list(models,
                                       NULL))
    for (x in models) {
      filename <- paste(modeldir,"/",
                        country,
                        "_ebma_",
                        "count",
                        "_",x,".RData",
                        sep = "")
      load(filename)
      for (i in 1:T) {
        predictions <- as.vector(ebma.results[[i]]$fit.oos)
        realizations <- as.vector(ebma.results[[i]]$actual.oos)
        # Cap predictions (some wild values distort the graph)
        predictions[predictions > max(realizations)] <- max(realizations)
        mse_dif <- (-(mean((realizations - predictions)^2)
                     - pop_auc_array[i]))
        if (country == 'indo') {
          auc_array[x, i] <- symSignRescale(mse_dif)
        } else {
          auc_array[x, i] <- mse_dif
        }
      }
    }
    mean_auc <- apply(auc_array, 1, mean)
    
    # Make Graph
    var.order1 <- order(mean_auc[names(figures_1_2_fixed)[-1]],
                        decreasing = FALSE)
    var.order2 <- order(mean_auc[names(figures_1_2_slow)],
                        decreasing = FALSE)
    var.order3 <- order(mean_auc[names(figures_1_2_annual)],
                        decreasing = FALSE)
    ranks1 = rep(NA,length(var.order1))
    ranks1[var.order1] <- ((1:length(var.order1))
                           +length(var.order2)
                           +length(var.order3)
                           +ifelse(length(var.order2)>0,
                                   2,1))
    ranks2 = rep(NA,length(var.order2))
    ranks2[var.order2] <- 1:length(var.order2)+length(var.order3)+1
    ranks3 = rep(NA,length(var.order3))
    ranks3[var.order3] <- 1:length(var.order3)
    ranks <- c(ranks1,ranks2,ranks3)
    varsets.total <- c(names(figures_1_2_fixed)[-1],
                       names(figures_1_2_slow),
                       names(figures_1_2_annual))
    pdf(paste("graphs/",
              country,
              "_AUCs_",
              "count",
              "_ebma_popplus.pdf",
              sep=""),
        width = 7, 
        height = 7)
    par(mar=c(5, 12, 2, 1) + 0.1)
    if (country=='colo') {
      xlim <- c(min(-5),5)
    } else {
      xlim <- c(-4.5,1.75)
    }
    plot(mean_auc[varsets.total],
         ranks,
         xlim = xlim,
         ylim = c(0,max(ranks)+ifelse(length(var.order2)>0,
                                      2,1)),
         ylab = "",
         xlab = "Performance improvement over population (AUC)",
         yaxt='n',
         xaxt='n',
         main = "")
    for (i in 1:length(ranks)) {
      lines(x = xlim,
            y = c(ranks[i],ranks[i]),
            col = "gray95")
    }
    points(as.vector(t(auc_array[varsets.total,])),
           rep(ranks,each = T),
           pch = rep(c(17,
                       rep(16,T-2),
                       15),
                     length(varsets.total)),
           cex = rep(c(1,
                       rep(.5,T-2),
                       1),
                     length(varsets.total)),
           col = rep(c("gray",
                       rep("black",T-2),
                       "gray"),
                     length(varsets.total)))
    points(mean_auc[varsets.total],
           ranks)
    axis(1,
         at = xtics[[country]]$ticks,
         labels = xtics[[country]]$labels,
         las = 1,
         cex.axis = .8)
    axis(2,
         at = ranks,
         labels = labels[varsets.total],
         las = 2)
    axis(2,
         at = c(length(ranks)+ifelse(length(var.order2)>0,
                                     3,2),
                length(ranks3)+1,
                length(ranks2)+length(ranks3)+ifelse(length(ranks2)>0,2,1)
         ),
         labels = c("Time-Invariant",
                    "Time varying (Annual)",
                    ifelse(length(ranks2)>0,"Slow-Moving","")),
         las = 2,
         font=4 )
    lines(x = c(0,0),
          y = c(-1,length(ranks)+ifelse(length(ranks2)>0,5,3)),
          lty = 3)
    dev.off()
}
